from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, update
from sqlalchemy.orm import Session
from app.models import IndicatorCodeCounter

def _to_base36(n: int, length: int) -> str:
    chars = '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ'
    s = ''
    if n == 0:
        s = '0'
    else:
        while n > 0:
            s = chars[n % 36] + s
            n //= 36
    return s.rjust(length, '0')

async def ensure_counter(session: AsyncSession, domain: str, length: int = 4) -> IndicatorCodeCounter:
    q = await session.execute(select(IndicatorCodeCounter).where(IndicatorCodeCounter.domain == domain))
    row = q.scalar_one_or_none()
    if row is None:
        row = IndicatorCodeCounter(domain=domain, current_value=0, base=36, code_length=length)
        session.add(row)
        await session.flush()
    return row

async def allocate_codes(session: AsyncSession, domain: str, count: int, length: int = 4) -> list[str]:
    q = select(IndicatorCodeCounter).where(IndicatorCodeCounter.domain == domain).with_for_update()
    res = await session.execute(q)
    row = res.scalar_one_or_none()
    if row is None:
        row = IndicatorCodeCounter(domain=domain, current_value=0, base=36, code_length=length)
        session.add(row)
        await session.flush()
    max_value = (row.base ** row.code_length) - 1
    start = row.current_value
    end = start + count
    if end > max_value:
        raise ValueError('code_exhausted')
    codes = [_to_base36(i, row.code_length) for i in range(start + 1, end + 1)]
    row.current_value = end
    await session.flush()
    return codes
